package com.eian.boot.common.core.utils;

import com.google.common.collect.Lists;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
 * @author alex.meng
 * @createTime 2025-09-12 14:26
 * @description 异步任务工具类
 */
@Slf4j
@Component
@AllArgsConstructor
public class AsyncTaskHelper {
    // 最大分片数量保护，防止内存溢出
    private static final int MAX_SHARDS = 200;
    private final Executor taskExecutor;

    /**
     * 执行一个无返回值的异步任务。
     *
     * @param task 要执行的任务 (Runnable)
     */
    public void run(Runnable task) {
        taskExecutor.execute(task);
    }

    /**
     * 执行一个有返回值的异步任务。
     *
     * @param task 要执行的任务 (Supplier<T>)
     * @param <T>  返回值类型
     * @return CompletableFuture<T> 用于获取异步计算结果或处理异常
     */
    public <T> CompletableFuture<T> run(Supplier<T> task) {
        return CompletableFuture.supplyAsync(task, taskExecutor);
    }

    /**
     * 执行分片任务 (无返回值收集)。
     * 将数据列表按指定大小分片，并行处理每个分片。
     * <p>
     * 调用者可以通过返回的 CompletableFuture 来：
     * - 调用 {@code future.join()} 阻塞等待所有分片完成。
     * - 调用 {@code future.thenRun(...)} 或 {@code future.whenComplete(...)}
     * 注册回调，以在所有任务完成后或发生异常时执行后续逻辑。
     * </p>
     *
     * @param dataList  需要处理的数据列表
     * @param shardSize 每个分片的大小
     * @param shardTask 处理单个分片的函数 (Consumer<List<T>>)
     * @param <T>       数据列表元素类型
     * @return CompletableFuture<Void> 在所有分片任务完成后完成，或在任一分片任务失败时失败。
     */
    public <T> CompletableFuture<Void> executeShardedTask(List<T> dataList, int shardSize, Consumer<List<T>> shardTask) {
        if (dataList == null || dataList.isEmpty()) {
            log.info("数据列表为空，直接返回已完成的 Future。");
            return CompletableFuture.completedFuture(null);
        }

        if (shardSize <= 0) {
            IllegalArgumentException e = new IllegalArgumentException("分片大小必须大于 0");
            log.warn("无效的分片大小：{}", shardSize);
            return CompletableFuture.failedFuture(e);
        }

        // 计算分片总数，防止恶意输入导致分片过多
        long shardCount = (dataList.size() + shardSize - 1) / shardSize;
        if (shardCount > MAX_SHARDS) {
            IllegalArgumentException e = new IllegalArgumentException(
                    "分片数量过多：" + shardCount + "，最大允许：" + MAX_SHARDS);
            log.warn("分片数量超出限制", e);
            return CompletableFuture.failedFuture(e);
        }

        List<CompletableFuture<Void>> futures;
        try {
            // 保留 Guava Lists.partition（兼容 Java 8~20）
            List<List<T>> partitions = Lists.partition(dataList, shardSize);

            futures = partitions.stream()
                    .map(shard -> CompletableFuture.runAsync(() -> {
                        try {
                            log.debug("正在处理大小为 {} 的分片", shard.size());
                            shardTask.accept(shard); // 执行分片任务
                            log.debug("分片处理完成，大小：{}", shard.size());
                        } catch (Exception e) {
                            log.error("处理分片时发生错误", e);
                            // 重新抛出运行时异常，使 CompletableFuture 失败
                            throw new RuntimeException(e);
                        }
                    }, taskExecutor))
                    .toList();
        } catch (Exception e) {
            log.error("分片或任务提交过程中发生错误", e);
            return CompletableFuture.failedFuture(e);
        }

        CompletableFuture<Void> allDoneFuture = CompletableFuture.allOf(futures.toArray(CompletableFuture[]::new));

        return allDoneFuture.whenComplete((result, throwable) -> {
            if (throwable != null) {
                log.info("一个或多个分片任务执行失败。");
            } else {
                log.info("所有分片任务已成功完成。");
            }
        });
    }

    /**
     * 执行分片任务并收集结果。
     * 将数据列表按指定大小分片，并行处理每个分片，并收集每个分片的处理结果。
     *
     * @param dataList  需要处理的数据列表
     * @param shardSize 每个分片的大小
     * @param shardTask 处理单个分片并返回结果的函数 (Function<List<T>, R>)
     * @param <T>       数据列表元素类型
     * @param <R>       分片任务返回结果类型
     * @return CompletableFuture<List < R>> 在所有分片任务完成后，
     * 包含所有分片结果的 CompletableFuture；若任一分片失败，则此 future 也会失败。
     */
    public <T, R> CompletableFuture<List<R>> executeShardedTaskWithResults(List<T> dataList, int shardSize, Function<List<T>, R> shardTask) {
        if (dataList == null || dataList.isEmpty()) {
            log.info("数据列表为空，返回空结果列表的 Future。");
            return CompletableFuture.completedFuture(List.of());
        }

        if (shardSize <= 0) {
            IllegalArgumentException e = new IllegalArgumentException("分片大小必须大于 0");
            log.warn("无效的分片大小：{}", shardSize);
            return CompletableFuture.failedFuture(e);
        }

        // 计算分片总数，防止恶意输入导致分片过多
        long shardCount = (dataList.size() + shardSize - 1) / shardSize;
        if (shardCount > MAX_SHARDS) {
            IllegalArgumentException e = new IllegalArgumentException(
                    "分片数量过多：" + shardCount + "，最大允许：" + MAX_SHARDS);
            log.warn("分片数量超出限制", e);
            return CompletableFuture.failedFuture(e);
        }

        List<CompletableFuture<R>> futures;
        try {
            List<List<T>> partitions = Lists.partition(dataList, shardSize);
            futures = partitions.stream()
                    .map(shard -> CompletableFuture.supplyAsync(() -> {
                        try {
                            log.debug("正在处理大小为 {} 的分片（需收集结果）", shard.size());
                            R result = shardTask.apply(shard);
                            log.debug("分片处理完成，大小：{}，结果类型：{}", shard.size(), result != null ? result.getClass().getSimpleName() : "null");
                            return result;
                        } catch (Exception e) {
                            log.error("处理分片（需收集结果）时发生错误", e);
                            throw new RuntimeException(e);
                        }
                    }, taskExecutor))
                    .toList();
        } catch (Exception e) {
            log.error("分片或任务提交过程中发生错误（需收集结果）", e);
            return CompletableFuture.failedFuture(e);
        }

        // 等待所有任务完成并收集结果
        return CompletableFuture.allOf(futures.toArray(CompletableFuture[]::new))
                .thenApply(v -> futures.stream()
                        .map(CompletableFuture::join)
                        .collect(Collectors.toList()));
    }
}